import argparse
import datetime

import tensorflow as tf


def create_model():
  return tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28), name='layers_flatten'),
    tf.keras.layers.Dense(512, activation='relu', name='layers_dense'),
    tf.keras.layers.Dropout(0.2, name='layers_dropout'),
    tf.keras.layers.Dense(10, activation='softmax', name='layers_dense_2')
  ])

def main():
  parser = argparse.ArgumentParser("Train MNIST using TensorFlow and TensorBoard")
  parser.add_argument('--epochs', type=int, default=5, help='训练轮数，越大，训练时间越久')
  parser.add_argument('--logdir', type=str, default="", help='TensorBoard 日志文件输出目录，默认为 logs/fit/<日期>')
  args = parser.parse_args()

  mnist = tf.keras.datasets.mnist
  (x_train, y_train),(x_test, y_test) = mnist.load_data()
  x_train, x_test = x_train / 255.0, x_test / 255.0

  model = create_model()
  model.compile(optimizer='adam',
                loss='sparse_categorical_crossentropy',
                metrics=['accuracy'])
  if args.logdir:
    logdir = args.logdir
  else:
    logdir = "logs/fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
  tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir, histogram_freq=1)

  # enable histogram computation every epoch with histogram_freq=1 (this is off by default)
  model.fit(x=x_train, 
            y=y_train, 
            epochs=args.epochs, 
            validation_data=(x_test, y_test), 
            callbacks=[tensorboard_callback])

if __name__ == '__main__':
  main()
